#!/usr/bin/env python2
#coding:utf-8

"""
Author:   Rusty Stiles --<rstiles74@gmail.com>
Purpose: Command line interface to rods_program
Created: 07/21/2011
Updated:
Comments:

"""

import socket
import sys, os, glob
import connection, plots, inputfile

class Menu:
    """Display a menu and respond to choices when run"""

    def __init__(self):
        self.ssh = connection
        self.user = os.getenv('USER', 'stiles')
        self.choices = {
            "1": self.change_user,
            "2": self.create_inputs,
            "4": self.start_runs,
            "0": self.quit_cli,
            "5": self.move_files,
            "3": self.compile_code,
            "6": self.get_progress,
            "7": self.get_ps,
            "8": self.run_command
            }

    def display_menu(self):
        print("Current user: " + self.user)
        print("""
            Menu:
            1. Change User
            2. Create Input Files
            3. Compile Code on Nodes
            4. Start Runs on the Nodes
            5. Move Files on the Nodes
            6. Check Progress of Running Code
            7. Check the Processor Status on the nodes
            8. Issue a Command on nodes
            0. Quit
            """)

    def run(self):
        while True:
            self.display_menu()
            choice = raw_input("\n\t    Enter an option: ").lower()
            print('\n\n')
            try:
                action = self.choices[str(choice)]
            except KeyError:
                print("\n\t    {} is not a valid option".format(choice))
            else:
                try:
                    action()
                except socket.error:
                    print("\n\n***Error connecting to host***\n\n")

    def change_user(self):
        self.user = raw_input("Username: ")

    def _get_matching_files(self, direc = 'cfgs', search = '*cfg'):
        folder = os.path.dirname(os.path.abspath(__file__))
        search = os.path.join(folder, direc, search)
        files = glob.glob(search)
        return files

    def _pick_file(self, files):
        cfg = None
        while not cfg:
            print("Pick an input file configuration")
            for key, val in enumerate(files):
                print(str(key) + ": " + os.path.basename(val))
            choice = raw_input("Pick the number corresponding to the correct file: ")
            try:
                cfg = files[int(choice)]
            except IndexError:
                print(choice + " is an invalid choice")
        return cfg

    def create_inputs(self):

        cfg_files = self._get_matching_files()
        if not cfg_files:
            print("No input file configuration found....")
            return
        cfg = self._pick_file(cfg_files)
        print(cfg)
        config = inputfile.readconfig(cfg)
        text = "Please enter the folder name where you want to store the input files."
        input_folder = self._make_folder(text)
        #config = self._edit_config(config)
        self._generate_inputs(input_folder, cfg, config)

    def _generate_inputs(self, input_folder, file_name, config):
        file_name = os.path.basename(file_name)
        baseName = file_name.replace(".cfg", "")
        num = None
        while not num:
            try:
                num = int(raw_input("How many input files would you like: "))
            except ValueError:
                print("That is not a number")
        baseFiles = [("c%.2i." % (i + 1) + baseName) for i in xrange(num)]
        files = [os.path.join(input_folder, fil) for fil in baseFiles]
        for fil in files:
            inputfile.writeinputs(config, fil)


    def get_path(self, text = "Enter folder name"):
        path = raw_input(text + 
                           "\n(Can be full path or relative path)\n" + ">> ")
        if path.startswith("~"):
            path = os.path.expanduser(path)
        else:
            path = os.path.abspath(path)
        return path

    def _make_folder(self, text):

        while True:
            folder_path = self.get_path(text)
            if os.path.exists(folder_path):
                print("Folder already exists, enter another folder or delete the folder and try again")
                continue
            else:
                print("Making Dir: " + folder_path)
                os.makedirs(folder_path)
                return folder_path


    def _edit_config(self, config):
        print("Not yet Implemented")
        return config

    def start_runs(self):
        print("Make sure nothing is running on the nodes you want to use and that you have retrieved " + 
              "all files that could be overwritten")
        text = "Enter the folder containing the input files"
        folder = self.get_path(text)
        nodes, work_folders = self.get_nodes_folders()
        files = glob.glob(os.path.join(folder, "c*"))
        number_files = len(files)
        nodes_count = len(nodes) * 2
        i = 0
        while (i < number_files) and (i < nodes_count):
            for node in nodes:
                client = self.ssh.Client(node, self.user)
                for folder in work_folders:
                    file_name = os.path.basename(files[i])
                    file_name = file_name[file_name.find(".") + 1 : ]
                    remote_file = os.path.join(folder, file_name)
                    client.put_file(files[i], remote_file)
                    i += 1
                client.close()
        self.ssh.start_runs_on_nodes(nodes, work_folders, username = self.user, timed = True)

    def get_nodes_folders(self, defaults = False):
        cluster, first, last = raw_input("Enter the Cluster ('e') first node and last node " + 
                                         "separated by commas e.g for nodes e65-e72 enter e,65,72:  ").split(",")
        cluster = cluster.lower()
        node_list = self.ssh.make_node_list(cluster, first, last)
        first_dir = '/work/' + self.user + '/'
        second_dir = '/work/' + self.user + '/b/'
        folders = [first_dir, second_dir]
        if not defaults:
            print("\ncurrent work folders are:\n " + first_dir + " and " + second_dir + 
                "\nInput different folders separated by a comma or hit enter to use defaults:   ")
            choice = raw_input().strip()
            if choice:
                folders = [folder for folder in choice.split(",") if folder]
        return node_list, folders

    def move_files(self):
        """****TODO Try ziping file together  for iteration and allowing multiple files input"""
        
        nodes, folders = self.get_nodes_folders()
        text = "What is the Name of the output file, e.g scdata.out   "
        output = raw_input(text)
        remote_files = [os.path.join(folder, output) for folder in folders]

        text = "Enter the path to the output folder to store results in"
        out_folder = self.get_path(text)
        num = len(nodes) * len(folders)
        outFiles = [("c%.2i." % (i + 1) + output) for i in xrange(num)]
        outPaths = [os.path.join(out_folder, outFile) for outFile in outFiles]
        i = 0
        for node in nodes:
            client = self.ssh.Client(node, self.user)
            print(node)
            for fil in remote_files:
                print(fil + "\t" + outPaths[i])
                client.move(fil, outPaths[i], delete_remote = True)
                i += 1

    def compile_code(self):
        nodes, work_folders = self.get_nodes_folders()
        print("Searching fortran_code folder for '.f' files")
        source_files = self._get_matching_files(direc = 'fortran_code', search = '*.f')
        source_file = self._pick_file(source_files)
        self.ssh.compile_for_nodes(nodes, source_file, work_folders, self.user)

    def run_command(self):
        nodes, folders = self.get_nodes_folders(defaults = True)
        text = "what command do you want to run on the nodes? "
        command = raw_input(text)
        self.ssh.run_commands(nodes, command, self.user)

    def get_progress(self):
        nodes, folders = self.get_nodes_folders()
        tempName = raw_input("Enter the name of the redirect file, e.g temp:   ")
        temp_files = [os.path.join(folder, tempName) for folder in folders]
        self.ssh.get_progress(nodes, temp_files, self.user)


    def get_ps(self):
        nodes, folders = self.get_nodes_folders(defaults = True)
        self.ssh.get_processor_status(nodes, self.user)


    def quit_cli(self):
        sys.exit(0)


if __name__ == '__main__':
    Menu().run()
